# -*- coding: utf-8 -*-
'''
Created on 25.10.2019

@author: yu03
'''

import numpy as np
import os
import re
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import glob
from mpl_toolkits.mplot3d import Axes3D
from Video_Unified import folder_path, np_result_names, hor_index, ver_index
from FFT_Interpolation import FFT_cal
from scipy.interpolate import interp1d
import sys

plt_show_mode = 'show' ### plt.show()
# plt_show_mode = 'save' ### plt.save()


def polyfit_1d(x, data, fs, degree, x_new_base):
    data_fit_params = np.polyfit(x, data, degree)
    data_fit_poly = np.poly1d(data_fit_params)
    data_fitline = data_fit_poly(x)
    nonlinearity = data - data_fitline
    
    length_fit_params = np.polyfit(x, x_new_base, 6)
    length_fit_poly = np.poly1d(length_fit_params)
    length_fitline = length_fit_poly(x)
    length_new = np.linspace(x_new_base[0], x_new_base[-1], num=len(x_new_base))
      
    spline = interp1d(length_fitline, nonlinearity, kind='cubic', bounds_error=False, fill_value=0)
    nonlinearity_P = spline(length_new)
    f, fft_xxx, FFT, phi = FFT_cal(nonlinearity, 1/fs)
    f_P, fft, FFT_P, phi = FFT_cal(nonlinearity_P, np.abs(length_new[0]-length_new[1]))
    
    return nonlinearity, length_new, nonlinearity_P, f[1:], FFT[1:], f_P[1:], FFT_P[1:]


''' 
    polyfit degree
'''     
length_dof = 10
angle_dof = 6
i = 0
fs = 47.6

''' 
    for all groups
''' 
for np_result in np_result_names:
    ''' 
        reading .npy file
    ''' 
    file_name = np_result.split('\\')[-1] ### x_lines.py
#     file_name = re.findall(r"openmode\\(.+).npy", np_result)[0]
    print(file_name)
    i += 1
    f = open(np_result, 'rb')
    f_size = os.fstat(f.fileno()).st_size
#     plt.savefig('C:/Users/yu03/Desktop/Gabor/Interferometer Sim/Simulation_Result/2. V_y_2 Modulation/Calc.Intensity_%i.jpg'%m, dpi=300)
    fig_name = folder_path + '\\' + file_name.split('_')[0] + '_FFT.png'

    
    ''' 
        put 4-D results in "lines"
    ''' 
    lines = []
    time_sequence = []
    while f.tell() < f_size:
        line = np.load(f, allow_pickle=True)
        lines.append(line)
    print('%i: %s, %eB'%(i, file_name, f_size))
    lines = np.array(lines)[::,::,300:2200]
    print(np.shape(lines), 'Channle, lines, frames:')
    frame_num = np.shape(lines)[2]
    line_num = np.shape(lines)[1]
    hor_lines = np.linspace(hor_index-(line_num-1)/2, hor_index+(line_num-1)/2, line_num, dtype='int')
    ver_lines = np.linspace(ver_index-(line_num-1)/2, ver_index+(line_num-1)/2, line_num, dtype='int') 
     
    hor_angle_set, ver_angle_set, hor_length_set, ver_length_set = lines
    ''' 
        making average
    ''' 
    averaged_results = []
     
    hor_angle_avr = np.mean(hor_angle_set, axis=0)
    ver_angle_avr = np.mean(ver_angle_set, axis=0)
    hor_length_avr = np.mean(hor_length_set, axis=0)
    ver_length_avr = np.mean(ver_length_set, axis=0)
     
    averaged_results.append(hor_angle_avr)
    averaged_results.append(ver_angle_avr)
    averaged_results.append(hor_length_avr)
    averaged_results.append(ver_length_avr)
     
 
    x = np.arange(frame_num)
     
    hor_angle_non, hor_angle_new, hor_angle_non_P, f_hor_angle_non, FFT_hor_angle_non, f_hor_angle_non_P, FFT_hor_angle_non_P = polyfit_1d(x, hor_angle_set[51], fs, angle_dof, x_new_base=hor_length_avr)
    hor_angle_avr_non, hor_angle_avr_new, hor_angle_avr_non_P, f_hor_angle_avr_non, FFT_hor_angle_avr_non, f_hor_angle_avr_non_P, FFT_hor_angle_avr_non_P = polyfit_1d(x, hor_angle_avr, fs, angle_dof, x_new_base=hor_length_avr)
    ver_angle_non, ver_angle_new, ver_angle_non_P, f_ver_angle_non, FFT_ver_angle_non, f_ver_angle_non_P, FFT_ver_angle_non_P = polyfit_1d(x, ver_angle_set[51], fs, angle_dof, x_new_base=ver_length_avr)
    ver_angle_avr_non, ver_angle_avr_new, ver_angle_avr_non_P, f_ver_angle_avr_non, FFT_ver_angle_avr_non, f_ver_angle_avr_non_P, FFT_ver_angle_avr_non_P = polyfit_1d(x, ver_angle_avr, fs, angle_dof, x_new_base=ver_length_avr)
    hor_length_non, hor_length_new, hor_length_non_P, f_hor_length_non, FFT_hor_length_non, f_hor_length_non_P, FFT_hor_length_non_P = polyfit_1d(x, hor_length_set[51], fs, length_dof, x_new_base=hor_length_avr)
    hor_length_avr_non, hor_length_avr_new, hor_length_avr_non_P, f_hor_length_avr_non, FFT_hor_length_avr_non, f_hor_length_avr_non_P, FFT_hor_length_avr_non_P = polyfit_1d(x, hor_length_avr, fs, length_dof, x_new_base=hor_length_avr)
    ver_length_non, ver_length_new, ver_length_non_P, f_ver_length_non, FFT_ver_length_non, f_ver_length_non_P, FFT_ver_length_non_P = polyfit_1d(x, ver_length_set[51], fs, length_dof, x_new_base=ver_length_avr)
    ver_length_avr_non, ver_length_avr_new, ver_length_avr_non_P, f_ver_length_avr_non, FFT_ver_length_avr_non, f_ver_length_avr_non_P, FFT_ver_length_avr_non_P = polyfit_1d(x, ver_length_avr, fs, length_dof, x_new_base=ver_length_avr)
 
       
    plt.figure('Nonlinearity FFT')
    plt.gcf().set_size_inches(18,9)
    
    plt.subplot(2,2,1)
    plt.loglog(1/f_hor_length_non_P, FFT_hor_angle_non, color='green', label='single-line-T')
    plt.loglog(1/f_hor_length_non_P, FFT_hor_angle_non_P, color='blue', label='single-line-P')
    plt.loglog(1/f_hor_length_avr_non_P, FFT_hor_angle_avr_non, color='black', label='avraged-T')
    plt.loglog(1/f_hor_length_avr_non_P, FFT_hor_angle_avr_non_P, color='red', label='avraged-P')
    plt.gca().invert_xaxis()
    plt.grid(which='both', axis='both')
    plt.title('Hor. Angle Nonlinearity FFT')
    plt.xlabel('Period / nm')
    plt.ylabel('Amplitude / urad')
    plt.xlim(4e3,30)
    plt.ylim(1e-4,10)
    plt.legend()
          
    plt.subplot(2,2,2)
    plt.loglog(1/f_hor_length_non_P, FFT_hor_length_non, color='green', label='single-line-T')
    plt.loglog(1/f_hor_length_non_P, FFT_hor_length_non_P, color='blue', label='single-line-P')
    plt.loglog(1/f_hor_length_avr_non_P, FFT_hor_length_avr_non, color='black', label='avraged-T')
    plt.loglog(1/f_hor_length_avr_non_P, FFT_hor_length_avr_non_P, color='red', label='avraged-P')
    plt.gca().invert_xaxis()
    plt.grid(which='both', axis='both')
    plt.title('Hor. Length Nonlinearity FFT')
    plt.xlabel('Period / nm')
    plt.ylabel('Amplitude / nm')
    plt.xlim(4e3,30)
    plt.ylim(1e-4,10)
    plt.legend()
    
    plt.subplot(2,2,3)
    plt.loglog(1/f_ver_length_non_P, FFT_ver_angle_non, color='green', label='single-line-T')
    plt.loglog(1/f_ver_length_non_P, FFT_ver_angle_non_P, color='blue', label='single-line-P')
    plt.loglog(1/f_ver_length_avr_non_P, FFT_ver_angle_avr_non, color='black', label='avraged-T')
    plt.loglog(1/f_ver_length_avr_non_P, FFT_ver_angle_avr_non_P, color='red', label='avraged-P')
    plt.gca().invert_xaxis()
    plt.grid(which='both', axis='both')
    plt.title('Ver. Angle Nonlinearity FFT')
    plt.xlabel('Period / nm')
    plt.ylabel('Amplitude / urad')
    plt.xlim(4e3,30)
    plt.ylim(1e-4,10)
    plt.legend()
      
    plt.subplot(2,2,4)
    plt.loglog(1/f_ver_length_non_P, FFT_ver_length_non, color='green', label='single-line-T')
    plt.loglog(1/f_ver_length_non_P, FFT_ver_length_non_P, color='blue', label='single-line-P')
    plt.loglog(1/f_ver_length_avr_non_P, FFT_ver_length_avr_non, color='black', label='avraged-T')
    plt.loglog(1/f_ver_length_avr_non_P, FFT_ver_length_avr_non_P, color='red', label='avraged-P')
    plt.gca().invert_xaxis()
    plt.grid(which='both', axis='both')
    plt.title('Ver. Length Nonlinearity FFT')
    plt.xlabel('Period / nm')
    plt.ylabel('Amplitude / nm')
    plt.xlim(4e3,30)
    plt.ylim(1e-4,10)
    plt.legend()
       
    #     figManager = plt.get_current_fig_manager()
    #     figManager.window.showMaximized()
    plt.tight_layout()
    if plt_show_mode == 'save':
        plt.savefig(fig_name, dpi=300)
        plt.close()
    elif plt_show_mode == 'show':
        plt.show()
    else:
        sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
    
    
plt.figure('Nonlinearity')
plt.gcf().set_size_inches(18,9)

plt.subplot(2,2,1)
# plt.plot(hor_angle_new, hor_angle_non, color='green', label='single-line-T')
# plt.plot(hor_angle_new, hor_angle_non_P, color='blue', label='single-line-P')
plt.plot(hor_angle_avr_new, hor_angle_avr_non, color='black', label='avraged-T')
plt.plot(hor_angle_avr_new, hor_angle_avr_non_P, color='red', label='avraged-P')
plt.grid(which='both', axis='both')
plt.title('Hor. Angle Nonlinearity FFT')
plt.xlabel('Pos. / nm')
plt.ylabel('Amplitude / urad')
plt.legend()

plt.subplot(2,2,2)
# plt.plot(hor_length_new, hor_length_non, color='green', label='single-line-T')
# plt.plot(hor_length_new, hor_length_non_P, color='blue', label='single-line-P')
plt.plot(hor_length_avr_new, hor_length_avr_non, color='black', label='avraged-T')
plt.plot(hor_length_avr_new, hor_length_avr_non_P, color='red', label='avraged-P')
plt.gca().invert_xaxis()
plt.grid(which='both', axis='both')
plt.title('Hor. Length Nonlinearity FFT')
plt.xlabel('Pos. / nm')
plt.ylabel('Amplitude / nm')
plt.legend()

plt.subplot(2,2,3)
# plt.plot(ver_angle_new, ver_angle_non, color='green', label='single-line-T')
# plt.plot(ver_angle_new, ver_angle_non_P, color='blue', label='single-line-P')
plt.plot(ver_angle_avr_new, ver_angle_avr_non, color='black', label='avraged-T')
plt.plot(ver_angle_avr_new, ver_angle_avr_non_P, color='red', label='avraged-P')
plt.grid(which='both', axis='both')
plt.title('Ver. Angle Nonlinearity FFT')
plt.xlabel('Pos. / nm')
plt.ylabel('Amplitude / urad')
plt.legend()
 
plt.subplot(2,2,4)
# plt.plot(ver_length_new, ver_length_non, color='green', label='single-line-T')
# plt.plot(ver_length_new, ver_length_non_P, color='blue', label='single-line-P')
plt.plot(ver_length_avr_new, ver_length_avr_non, color='black', label='avraged-T')
plt.plot(ver_length_avr_new, ver_length_avr_non_P, color='red', label='avraged-P')
plt.gca().invert_xaxis()
plt.grid(which='both', axis='both')
plt.title('Ver. Length Nonlinearity FFT')
plt.xlabel('Pos. / nm')
plt.ylabel('Amplitude / nm')
plt.legend()
 
 
# figManager = plt.get_current_fig_manager()
# figManager.window.showMaximized()
plt.tight_layout()

if plt_show_mode == 'save':
    plt.savefig(folder_path + '\\' + 'Nonlinearity.png', dpi=300)
    plt.close()
elif plt_show_mode == 'show':
    plt.show()
else:
    sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)
    
plt.figure('Raw Data')
plt.gcf().set_size_inches(18,9)

plt.subplot(2,1,1)
plt.plot(hor_angle_avr, color='blue', label='hor')
plt.plot(ver_angle_avr, color='red', label='red')
plt.grid(which='both', axis='both')
plt.title('Angle')
#     plt.xlabel('Period / nm')
plt.ylabel('Angle / urad')
plt.legend()

plt.subplot(2,1,2)
plt.plot(hor_length_avr, color='blue', label='hor')
plt.plot(ver_length_avr, color='red', label='red')
plt.grid(which='both', axis='both')
plt.title('Length')
#     plt.xlabel('Period / nm')
plt.ylabel('Length / nm')
plt.legend()

#     figManager = plt.get_current_fig_manager()
#     figManager.window.showMaximized()
plt.tight_layout()

if plt_show_mode == 'save':
    plt.savefig(folder_path + '\\' + 'Raw_Data.png', dpi=300)
    plt.close()
elif plt_show_mode == 'show':
    plt.show()
else:
    sys.exit('Figure Save/Show Error:\n plt_show_mode = %s'%plt_show_mode)